package com.pkk.cloud.support.gateway.limiter.limit.aop;

import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.common.util.concurrent.RateLimiter;
import com.pkk.cloud.support.gateway.code.GatewayBusinessCode;
import com.pkk.cloud.support.gateway.limiter.limit.EnableRateLimit;
import com.pkk.cloud.support.gateway.limiter.limit.annotations.RateLimit;
import com.pkk.components.rpc.response.util.ResponseUtil;
import java.io.IOException;
import java.lang.reflect.Method;
import java.util.concurrent.ConcurrentHashMap;
import javax.servlet.ServletOutputStream;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.Signature;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.autoconfigure.condition.ConditionalOnBean;
import org.springframework.context.annotation.Configuration;
import org.springframework.http.MediaType;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

/**
 * @description: 接口的限制aop
 * @author: peikunkun
 * @create: 2020-03-24 14:05
 **/
@Slf4j
@Aspect
@Configuration
@ConditionalOnBean(EnableRateLimit.class)
public class RateLimitAspect {

  @Autowired
  private HttpServletResponse response;
  //用来存放不同接口的RateLimiter(key为接口名称，value为RateLimiter)
  //在有RBAC框架的项目中，这个地方可以使用加载数据库配置的办法
  private ConcurrentHashMap<String, RateLimiter> map = new ConcurrentHashMap<>();
  private static ObjectMapper objectMapper = new ObjectMapper();

  static {
    objectMapper.setSerializationInclusion(JsonInclude.Include.NON_NULL);
  }


  /**
   * @Description: 增加切面设置接口拦截
   * @Param: joinPoint
   * @Param rateLimit
   * @return: java.lang.Object
   * @Author: peikunkun
   * @Date: 2020/3/24 0024 下午 2:09
   */
  @Around(value = "@annotation(rateLimit)")
  public Object around(ProceedingJoinPoint joinPoint, RateLimit rateLimit) throws NoSuchMethodException {
    final HttpServletRequest request = ((ServletRequestAttributes) RequestContextHolder.getRequestAttributes())
        .getRequest();
    final String remoteAddr = request.getRemoteAddr();

    Object obj = null;
    //获取拦截的方法名
    Signature sig = joinPoint.getSignature();
    MethodSignature msig = (MethodSignature) sig;
    //返回被织入增加处理目标对象
    Object target = joinPoint.getTarget();
    //为了获取注解信息
    Method currentMethod = target.getClass().getMethod(msig.getName(), msig.getParameterTypes());
    //获取注解信息
    RateLimit annotation = currentMethod.getAnnotation(RateLimit.class);
    double limitNum = annotation.limitNum(); //获取注解每秒加入桶中的token
    // 注解所在类名加方法名区分不同的限流策略
    String functionName = remoteAddr.concat(target.getClass().getSimpleName()).concat(msig.getName());

    //获取rateLimiter
    RateLimiter rateLimiter = null;
    if (map.containsKey(functionName)) {
      rateLimiter = map.get(functionName);
    } else {
      map.put(functionName, RateLimiter.create(limitNum));
      rateLimiter = map.get(functionName);
    }

    try {
      if (rateLimiter.tryAcquire()) {
        //执行方法
        obj = joinPoint.proceed();
      } else {
        //拒绝了请求（服务降级）
        String result = objectMapper.writeValueAsString(ResponseUtil.error(GatewayBusinessCode.TOO_MANY_REQUESTS));
        outErrorResult(result);
      }
    } catch (Throwable throwable) {
      throwable.printStackTrace();
    }
    return obj;
  }


  /**
   * @Description: 将结果返回
   * @Param: result
   * @return: void
   * @Author: peikunkun
   * @Date: 2020/3/24 0024 下午 2:17
   */
  public void outErrorResult(String result) {
    response.setContentType(MediaType.APPLICATION_JSON_UTF8_VALUE);
    try (ServletOutputStream outputStream = response.getOutputStream()) {
      outputStream.write(result.getBytes("utf-8"));
    } catch (IOException e) {
      e.printStackTrace();
    }
  }
}
